package plan

import (
	"mini-sql/expression"
	"mini-sql/model"
)

type SelectPlanNode struct {
	source          PlanNode
	selectStatement *model.SelectStatement
	fieldNames      []string
}

// Prepare implements PlanNode.
func (s *SelectPlanNode) Prepare() error {
	fieldNames := make([]string, 0)

	for _, expr := range s.selectStatement.ColumnNames {
		name, err := expression.GetFieldName(expr)
		if err != nil {
			return err
		}
		fieldNames = append(fieldNames, name)
	}

	s.fieldNames = fieldNames
	return s.source.Prepare()
}

// GetNextTuple implements PlanNode.
func (s *SelectPlanNode) GetNextTuple() (*model.TableTuple, error) {
	tuple, err := s.source.GetNextTuple()
	if err != nil {
		return nil, err
	}
	if tuple == nil {
		return nil, nil
	}
	return s.getSelectedColumn(tuple)
}

func (s *SelectPlanNode) getSelectedColumn(tuple *model.TableTuple) (*model.TableTuple, error) {
	columns := s.fieldNames
	if len(columns) == 1 && columns[0] == "*" {
		return tuple, nil
	}
	newFields := make([]model.TableField, 0)
	for _, fieldName := range columns {
		value := tuple.GetValueByName(fieldName)
		if value == nil {
			value = &model.TableField{
				Name:  fieldName,
				Value: nil,
			}
		}
		newFields = append(newFields, *value)
	}
	return &model.TableTuple{
		RowID:  tuple.RowID,
		Fields: newFields,
	}, nil
}

func NewSelectPlanNode(source PlanNode, selectStatement *model.SelectStatement) PlanNode {
	return &SelectPlanNode{
		source:          source,
		selectStatement: selectStatement,
	}
}
